//
// This file is part of cpp-lib, a library of C++ functions and classes.
// Cpp-lib is copyright (C) 2004 and onwards, The CPL Interest Group.
// 
// This library is free software; you can redistribute it and/or modify it
// under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation; either version 2.1 of the License, or (at
// your option) any later version.
// 
// This library is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser
// General Public License for more details.
// 
// You should have received a copy of the GNU Lesser General Public License
// along with this library (see the file COPYING); if not, write to the
// Free Software Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
// 02111-1307  USA
//
// $Id$
//

#ifndef CPP_LIB_CRYPT_STREAM_H
#define CPP_LIB_CRYPT_STREAM_H

#include <fstream>
#include <algorithm>
#include <ios>
#include <memory>

#include <cstring>
#include <cstdlib>
#include <cstdio>
#include <cassert>

#include "cpp-lib/util.h"


namespace cpl {

namespace crypt    {


//
// Output and input stream buffers implementing CFB for a block cipher CB.
//
// Algorithm:  See Menenzes, Handbook of Applied Cryptography, Algorithm
// 7.17 (p. 231).
//


//
// The encrypting output buffer.
//

template< typename BC >
struct cfb_ostreambuf : std::streambuf {

  //
  // Set the cipher and initialization vector to use.
  //

  cfb_ostreambuf( 
    BC const& c ,
    std::vector< char > const& IV ,
    std::auto_ptr< std::streambuf > sink  
  ) 
  : cipher( cipher ) ,
    I( IV ) ,
    sink( sink )
  { 

    if( IV.size() != c.block_size() ) 
    { throw std::runtime_error
      ( "cfb: bad initialization vector size" ) ; }

  }

protected:

  //
  // Convention:  x is plaintext, c is ciphertext.
  //

  virtual int_type overflow( int_type x ) {

    // debug version:  one-to-one mapping.
    // return sink->sputc( x ) ;

    if( traits_type::eof() == x ) { return traits_type::eof() ; }

    std::vector< char > O = I ;
    cipher.encrypt_block( O ) ;

    char const t = O[ 0 ] ;
    char const c = x ^ t ;

    // Shift I one char to the right.
    std::copy_backward( I.begin() , I.end() - 1 , I.end() ) ;
    I[ 0 ] = c ;

    return sink->sputc( c ) ;

  }

  virtual int sync() { return sink->pubsync() ; }

private:

  BC cipher             ;
  std::vector< char > I ;
  std::auto_ptr< std::streambuf > const sink ;

} ;


// 
// The decrypting input buffer.
//

template< typename BC >
struct cfb_istreambuf : std::streambuf {

  //
  // Set the cipher and initialization vector to use.
  //

  cfb_istreambuf( 
    BC const& c ,
    std::vector< char > const& IV ,
    std::auto_ptr< std::streambuf > source 
  ) 
  : cipher( cipher ) ,
    I     ( IV     ) ,
    source( source )
  { 

    setg( 
      buffer + size_pb ,     // beginning of putback area
      buffer + size_pb ,     // read position
      buffer + size_pb       // end position
    ) ;   

    if( IV.size() != c.block_size() ) 
    { throw std::runtime_error
      ( "cfb: bad initialization vector size" ) ; }

  }


protected:

  virtual int_type underflow () {

    // is read position before end of buffer?

    if( gptr() < egptr() ) 
    { return traits_type::to_int_type( *gptr() ) ; }

    /* process size of putback area
     * - use number of characters read
     * - but at most size of putback area
     */
    int const n_pb = 
      std::min( 
        static_cast< long >( gptr() - eback() ) , 
        static_cast< long >( size_pb          ) 
      ) ;

    /* copy up to size_pb characters previously read into
     * the putback area
     */
    std::memmove( buffer + size_pb - n_pb , gptr() - n_pb , n_pb ) ;

    // Read at most buf_size new characters and put them beginning
    // from buffer + size_pb.
    int i ;
    for( i = 0 ; i < buf_size ; ++i ) {

      int_type const c = source->sbumpc() ;

      if( traits_type::eof() == c ) { break ; }

      std::vector< char > O = I ;
#if 0 
      std::cerr 
        << int( O[ 0 ] ) << ' ' << int( O[ 1 ] ) << ' '
        << int( O[ 2 ] ) << ' ' << int( O[ 3 ] ) << ' '
        << int( O[ 4 ] ) << ' ' << int( O[ 5 ] ) << ' '
        << int( O[ 6 ] ) << ' ' << int( O[ 7 ] )
        << '\n' 
      ;
#endif

      cipher.encrypt_block( O ) ;

#if 0
      std::cerr 
        << int( O[ 0 ] ) << ' ' << int( O[ 1 ] ) << ' '
        << int( O[ 2 ] ) << ' ' << int( O[ 3 ] ) << ' '
        << int( O[ 4 ] ) << ' ' << int( O[ 5 ] ) << ' '
        << int( O[ 6 ] ) << ' ' << int( O[ 7 ] )
        << '\n' 
      ;
#endif


      char const t = O[ 0 ] ;
      char const x = traits_type::to_char_type( c ) ^ t ;

      // Shift I one char to the right.
      std::copy_backward( I.begin() , I.end() - 1 , I.end() ) ;
      I[ 0 ] = c ;

      *( buffer + size_pb + i ) = x ;
      // *( buffer + size_pb + i ) = traits_type::to_char_type( c ) ;

    }

    if( 0 == i ) { return traits_type::eof() ; }

    // reset buffer pointers
    setg(
      buffer + size_pb - n_pb , // beginning of putback area
      buffer + size_pb        , // read position
      buffer + size_pb + i      // end of buffer
    ) ;

    return traits_type::to_int_type( *gptr() ) ;

  }

private:

  // 
  // Data buffer:
  //
  // - at most, size_pb characters in putback area plus
  // - at most, buf_size characters in ordinary read buffer
  // 
  
  enum { size_pb  = 4    } ;  // size of putback area
  enum { buf_size = 1024 } ;  // size of the data buffer

  char buffer[ buf_size + size_pb ] ;  // data buffer

  BC cipher              ;
  std::vector< char > I  ;
  std::auto_ptr< std::streambuf > const source ;

} ;



//
// Open a file for reading while decrypting it using the given block 
// cipher in CFB mode with the given initialization vector.
//
// The file is opened by delegation to open_readbuf().
//

template< typename BC >
std::auto_ptr< std::istream > open_read(
  BC const& cipher              ,
  std::vector< char > const& IV ,
  std::string const& name       ,
  std::string      & which      ,
  std::vector< std::string > const& path = std::vector< std::string >()
) {

  std::auto_ptr< std::streambuf > buf ;
  buf.reset( cpl::util::file::open_readbuf( name , which , path ).release() ) ; 

  return std::auto_ptr< std::istream >( 
    new cpl::util::file::owning_istream( 
      std::auto_ptr< std::streambuf >( 
        new cfb_istreambuf< BC >( cipher , IV , buf ) 
      )
    ) 
  ) ;

}


//
// Equivalent to the above function but without information about
// the opened file.
//

template< typename BC >
std::auto_ptr< std::istream > open_read(
  BC const& cipher              ,
  std::vector< char > const& IV ,
  std::string const& name  ,
  std::vector< std::string > const& path = std::vector< std::string >()
) { 
  
  std::string dummy ; 
  return open_read( cipher , IV , name , dummy , path ) ; 

}


//
// Open a file for writing while encrypting it using the given block 
// cipher in CFB mode with the given initialization vector.
//
// The file is opened by delegation to open_writebuf().
//

template< typename BC >
std::auto_ptr< std::ostream > open_write(
  BC const& cipher              ,
  std::vector< char > const& IV ,
  std::string const& name 
) {
  
  std::auto_ptr< std::streambuf > buf ;
  buf.reset( cpl::util::file::open_writebuf( name ).release() ) ;

  return std::auto_ptr< std::ostream >( 
    new cpl::util::file::owning_ostream( 
      std::auto_ptr< std::streambuf >( 
        new cfb_ostreambuf< BC >( cipher , IV , buf ) 
      ) 
    ) 
  ) ;

}


} // namespace crypt

} // namespace cpl


#endif // CPP_LIB_CRYPT_STREAM_H
